#pragma once

#include <atlcoll.h>

namespace Common
{
namespace Windows
{
namespace Threading
{
	typedef VOID (CALLBACK *POOL_CALLBACK)(PVOID Context);

	class CPoolContextBase
	{
	public:
		CPoolContextBase(POOL_CALLBACK callback, PVOID pContext= NULL) : m_pContext(pContext), m_callback(callback)
		{

		}
		PVOID GetContext()
		{
			return m_pContext;
		}

		POOL_CALLBACK GetCallback()
		{
			return m_callback;
		}
	protected:
		virtual PVOID GetPoolContext() = 0;
	private:		
		PVOID m_pContext;
		POOL_CALLBACK m_callback;
	};
	
	class CPoolContextVista : public CPoolContextBase
	{
	public:
		CPoolContextVista(POOL_CALLBACK callback, PVOID pContext= NULL) : CPoolContextBase(callback, pContext)
		{

		}

		CPoolContextVista() : CPoolContextBase(NULL, NULL)
		{

		}

		void SetWorkObject(PTP_WORK pWork)
		{
			m_pWork = pWork;
		}
	protected:
		PVOID GetPoolContext()
		{
			return m_pWork;
		}
	private:
		PTP_WORK m_pWork;
	};

class CVistaThreadPool
{
public:
	CVistaThreadPool(bool defaultPool = true) : m_pPool(NULL), m_dwPoolMin(0), m_dwPoolMax(0), m_bDefaultInit(defaultPool)
	{
	
	}
	

	~CVistaThreadPool()
	{
		if (NULL != m_pPool)
		{
			::CloseThreadpool(m_pPool);
		}
	}	
	
	bool SetPoolThreadMimimum(DWORD dwMin)
	{
		if (dwMin != m_dwPoolMin && 0 < dwMin)
		{
			if (::SetThreadpoolThreadMinimum(GetPool(), dwMin))
				m_dwPoolMin = dwMin;
			else
				return false;
		}
		return true;
	}
	void SetPoolThreadMaximum(DWORD dwMax)
	{
		if (dwMax != m_dwPoolMax && 0 < dwMax)
		{
			m_dwPoolMax = dwMax;
			::SetThreadpoolThreadMaximum(GetPool(), dwMax);			
		}
		
	}
	CPoolContextBase * CreateWorkObject (POOL_CALLBACK callback, PVOID pData = NULL)
	{
		if (NULL == callback)
		{
			return NULL;
		}
		CPoolContextVista data(callback, pData);
		
		PTP_WORK pWorkObj = ::CreateThreadpoolWork(WrappertCallback, this, &this->GetCallbackEnv());
		data.SetWorkObject(pWorkObj);
		m_arrWorkObjects.SetAt(pWorkObj, data);
		
		return &m_arrWorkObjects[pWorkObj];
	}
	void SubmitWorkItem(PVOID pWorkObject)
	{
		if (NULL == pWorkObject)
		{
			return ;
		}

		::SubmitThreadpoolWork((PTP_WORK)pWorkObject);
	}

protected:

	PTP_POOL GetPool() {
		if (NULL == m_pPool)
		{
			m_pPool = ::CreateThreadpool(NULL);
			::InitializeThreadpoolEnvironment(&m_poolCBEnviron);
			// Associate the thread pool with the calback environment.
			::SetThreadpoolCallbackPool(&m_poolCBEnviron, m_pPool);
			if (m_bDefaultInit)
			{
				InitDefaultPool();
			}
		}
		return m_pPool;
	}



	TP_CALLBACK_ENVIRON & GetCallbackEnv()
	{
		if (NULL == m_pPool)
		{
			this->GetPool();
		}
		return m_poolCBEnviron;
	}

	void InitDefaultPool()
	{
		SYSTEM_INFO info = {0};
		::GetNativeSystemInfo(&info);
		DWORD dwSize = 2 < info.dwNumberOfProcessors ? 2 : info.dwNumberOfProcessors;

		this->SetPoolThreadMimimum(dwSize);
		this->SetPoolThreadMaximum(dwSize);

	}
	
	static VOID CALLBACK WrappertCallback(PTP_CALLBACK_INSTANCE Instance,
								  PVOID pContext,
								  PTP_WORK Work)
	{
		CVistaThreadPool * pool = reinterpret_cast<CVistaThreadPool *>(pContext);
	}

private:
	PTP_POOL m_pPool;
	DWORD m_dwPoolMin;
	DWORD m_dwPoolMax;
	TP_CALLBACK_ENVIRON m_poolCBEnviron;
	CAtlMap<PTP_WORK, CPoolContextVista> m_arrWorkObjects;
	bool m_bDefaultInit;
};

} //namespace Threading
} //namespace Windows
} //namespace Threading